import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error,mean_absolute_error
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv1D, MaxPooling1D, LSTM, Dense, Flatten, Reshape
import matplotlib.pyplot as plt


# 数据加载和预处理
def load_and_preprocess_data(file_path, input_time_steps, output_time_steps):
    df = pd.read_csv(file_path)
    scaler = MinMaxScaler()
    scaled_data = scaler.fit_transform(df)
    num_features = scaled_data.shape[1]

    X, y = [], []
    for i in range(len(scaled_data) - input_time_steps - output_time_steps + 1):
        X.append(scaled_data[i:(i + input_time_steps)])
        y.append(scaled_data[(i + input_time_steps):(i + input_time_steps + output_time_steps)])

    X = np.array(X)
    y = np.array(y)

    split_index = int(len(X) * 0.8)
    X_train, X_test = X[:split_index], X[split_index:]
    y_train, y_test = y[:split_index], y[split_index:]

    return X_train, X_test, y_train, y_test, scaler, num_features


# 构建和训练 CNN-LSTM 模型
def cnn_lstm_model(X_train, y_train, input_time_steps, num_features, output_time_steps, epochs):
    model = Sequential([
        Conv1D(filters=64, kernel_size=3, activation='relu', input_shape=(input_time_steps, num_features)),
        MaxPooling1D(pool_size=2),
        LSTM(units=64, return_sequences=True),
        LSTM(units=32, return_sequences=True),
        Flatten(),
        Dense(units=output_time_steps * num_features, activation='linear'),
        Reshape((output_time_steps, num_features))
    ])

    model.compile(optimizer='adam', loss='mean_squared_error')
    model.fit(X_train, y_train, epochs=epochs, batch_size=32)
    return model


# 生成预测并反向缩放
def generate_predictions(model, X_test, y_test, scaler, num_features):
    y_pred = model.predict(X_test)
    y_test_actual = scaler.inverse_transform(y_test.reshape(-1, num_features))
    y_pred_actual = scaler.inverse_transform(y_pred.reshape(-1, num_features))
    return y_test_actual, y_pred_actual


# 可视化预测值和真实值的对比
def plot_predictions(y_test_actual, y_pred_actual):
    time_index = range(len(y_test_actual))

    plt.figure(figsize=(12, 6))
    plt.plot(time_index, y_test_actual[:, 0], label='Actual', color='blue')
    plt.plot(time_index, y_pred_actual[:, 0], label='Predicted', color='red', linestyle='--')
    plt.xlabel('Time Steps')
    plt.ylabel('Values')
    plt.title('Actual vs Predicted Values')
    plt.legend()
    plt.grid(True)
    plt.show()


# 计算均方误差
def calculate_mse(y_test_actual, y_pred_actual):
    mse = mean_squared_error(y_test_actual, y_pred_actual)
    print(f"Mean Squared Error: {mse}")
    return mse


# 主程序
if __name__ == '__main__':
    input_time_steps = 10
    output_time_steps = 10
    epochs = 50
    file_path = 'datasets/kalman_data.csv'

    # 加载和预处理数据
    X_train, X_test, y_train, y_test, scaler, num_features = load_and_preprocess_data(file_path, input_time_steps,
                                                                                      output_time_steps)

    # 构建并训练模型
    model = cnn_lstm_model(X_train, y_train, input_time_steps, num_features, output_time_steps, epochs)

    # 生成预测并计算误差
    y_test_actual, y_pred_actual = generate_predictions(model, X_test, y_test, scaler, num_features)
    plot_predictions(y_test_actual, y_pred_actual)
    mse = calculate_mse(y_test_actual, y_pred_actual)
    mae = mean_absolute_error(y_test_actual, y_pred_actual)
    print(f"Mean Absolute Error: {mae}")